SDFT: Self-Distillation Enables Continual Learning

1. Challenge
2. SDFT Solution
- 2.1. From The Perspective of Inverse RL
- 2.2. In-Context Learning Assumption
  - 2.2.1. Empirical Validation of Assumption
3. Experiments
4. Ablation Studies
- 4.1. What to Demonstrate?

The paper proposes Self-Distillation Fine-Tuning (SDFT) that aims to handle the problem of continual learning: keep learning new skills while not forgetting old skills.

1. Challenge

2. SDFT Solution

SDFT uses same model as both teacher and student. Given a foundation model with policy \(\pi\), the teacher is conditioned with task description \(x\) and expert demonstration \(c\), denoted by \(\pi(\cdot | x,c)\); while the student is solely conditioned by task description \(x\), denoted by \(\pi_{\theta}(\cdot | x)\)

For every prompt \(x\), SDFT samples responses from the student policy \(y \sim \pi_{\theta}(\cdot | x)\) and minimizes the reverse KL divergence between the student and the teacher distributions:

\begin{equation} \mathcal{L}(\theta) = D_{KL}(\pi_{\theta}(\cdot | x) \;\|\; \pi(\cdot | x,c)) = \mathbb{E}_{y\sim\pi_{\theta}(y\mid x)}\left[ \log\frac{\pi_{\theta}(y | x)}{\pi(y | x,c)} \right] \end{equation}

Since LLMs are autoregressive, we can rewrite the objective into token-level loss. Then, we can represent the gradient w.r.t. the student parameters \(\theta\) while treating the teacher distributions as fixed. We can get the following gradient estimator.

\begin{equation} \nabla_{\theta}\mathcal{L}(\theta) = \mathbb{E}_{y\sim\pi_{\theta}}\left[ \sum_{t}\sum_{y_{t} \in \mathcal{V}}\log\frac{\pi_{\theta}(y_{t} | y_{\lt t}, x)}{\pi(y_{t} | y_{\lt t}, x, c)} \nabla_{\theta}\log\pi_{\theta}(y_{t} | y_{\lt t}, x) \right] \end{equation}

For the teacher model, SDFT uses Exponential Moving Average (EMA) of the student parameters.

2.1. From The Perspective of Inverse RL

2.2. In-Context Learning Assumption

The ICL Assumption tells that, the optimal policy for the task can be approximated by \(\pi^{\ast}_{k+1}(y | x) \approx \pi(y | x,c)\) and therefore is a good teacher. The quality of this approximation depends on 2 conditions:

Optimality.
Minimal deviation.

2.2.1. Empirical Validation of Assumption

3. Experiments

3.1. Effect of Model Size

Model size may influence model’s in-context learning ability, thus influencing performance of SDFT. Larger model size results in better in-context learning ability.

3.2. Training Reasoning Models w/o Reasoning Data

3.3. What Drives the Improvement in Performance?

Paper tested 3 variants of student models for evaluating:

student model with SDFT
SFT from teacher model, where student is trained offline on generated trajectories
offline distillation from teacher model, where student is trained to minimize the KL loss on a fixed dataset of teacher generated outputs.

The results shows that SDFT always outperforms SFT and offline distillation.